import os
import json

# 输入目录
input_dir = 'dataset/houronghe/houronghe1/annotations'
# 输出目录
output_dir = 'dataset/houronghe/houronghe1/label'
os.makedirs(output_dir, exist_ok=True)

# label_id 与 obj_type 映射
label_map = {
    0: "forklift",
    1: "tractor",
    2: "water_horse",
    3: "flatbed",
    4: "guardrail",
    5: "truck",
    6: "car",
    7: "bicycle",
    8: "pedestrian",
    9: "cyclist",
    10: "trailer",
    11: "traffic_cone",
    12: "unknown"
}


def convert_annotation_to_single_frame(annotation, label_id):
    obj_type = label_map.get(label_id, "unknown")
    position = annotation["position"]
    rotation = annotation["rotation"]
    scale = annotation["scale"]

    return {
        "obj_id": "1",
        "obj_type": obj_type,
        "psr": {
            "position": {
                "x": position[0],
                "y": position[1],
                "z": position[2]
            },
            "rotation": {
                "x": rotation[0],
                "y": rotation[1],
                "z": rotation[2]
            },
            "scale": {
                "x": scale[0],  
                "y": scale[1],  
                "z": scale[2]
            }
        }
    }

# 遍历 input_dir 下所有 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith('.json'):
        input_path = os.path.join(input_dir, filename)
        with open(input_path, 'r') as f:
            data = json.load(f)

        for item in data.get("items", []):
            frame_name = item["point_cloud"]["path"]
            annotations = item.get("annotations", [])
            objs = []

            for ann in annotations:
                label_id = ann["label_id"]
                obj = convert_annotation_to_single_frame(ann, label_id)
                objs.append(obj)

            output_path = os.path.join(output_dir, frame_name.replace(".pcd", ".json"))
            with open(output_path, "w") as out_f:
                json.dump(objs, out_f, indent=2)

print("The label files have been generated successfully.")
